module camsrfexch

  !-----------------------------------------------------------------------
  ! Module to handle data that is exchanged between the CAM atmosphere
  ! model and the surface models (land, sea-ice, and ocean).
  !-----------------------------------------------------------------------

  use shr_kind_mod,    only: r8 => shr_kind_r8, r4 => shr_kind_r4
  use constituents,    only: pcnst
  use ppgrid,          only: pcols, begchunk, endchunk
  use phys_grid,       only: get_ncols_p, phys_grid_initialized
  use infnan,          only: posinf, assignment(=)
  use cam_abortutils,  only: endrun
  use cam_logfile,     only: iulog
  use srf_field_check, only: active_Sl_ram1, active_Sl_fv, active_Sl_soilw,                &
                             active_Fall_flxdst1, active_Fall_flxvoc, active_Fall_flxfire, &
                             active_Faxa_nhx, active_Faxa_noy

  implicit none
  private

  public atm2hub_alloc              ! Atmosphere to surface data allocation method
  public hub2atm_alloc              ! Merged hub surface to atmosphere data allocation method
  public atm2hub_deallocate
  public hub2atm_deallocate
  public cam_export
  public cam_out_t                  ! Data from atmosphere
  public cam_in_t                   ! Merged surface data

  !---------------------------------------------------------------------------
  ! This is the data that is sent from the atmosphere to the surface models
  !---------------------------------------------------------------------------

  type cam_out_t
    integer  :: lchnk                        ! chunk index
    integer  :: ncol                         ! number of columns in chunk
    real(r8) :: tbot    (pcols)              ! bot level temperature
    real(r8) :: zbot    (pcols)              ! bot level height above surface
    real(r8) :: topo    (pcols)              ! surface topographic height (m)
    real(r8) :: ubot    (pcols)              ! bot level u wind
    real(r8) :: vbot    (pcols)              ! bot level v wind
    real(r8) :: qbot    (pcols,PCNST)        ! bot level specific humidity
    real(r8) :: pbot    (pcols)              ! bot level pressure
    real(r8) :: rho     (pcols)              ! bot level density
    real(r8) :: netsw   (pcols)              !
    real(r8) :: flwds   (pcols)              !
    real(r8) :: precsc  (pcols)              !
    real(r8) :: precsl  (pcols)              !
    real(r8) :: precc   (pcols)              !
    real(r8) :: precl   (pcols)              !
    real(r8) :: soll    (pcols)              !
    real(r8) :: sols    (pcols)              !
    real(r8) :: solld   (pcols)              !
    real(r8) :: solsd   (pcols)              !
    real(r8) :: thbot   (pcols)              !
    real(r8) :: co2prog (pcols)              ! prognostic co2
    real(r8) :: co2diag (pcols)              ! diagnostic co2
    real(r8) :: psl     (pcols)
    real(r8) :: bcphiwet(pcols)              ! wet deposition of hydrophilic black carbon
    real(r8) :: bcphidry(pcols)              ! dry deposition of hydrophilic black carbon
    real(r8) :: bcphodry(pcols)              ! dry deposition of hydrophobic black carbon
    real(r8) :: ocphiwet(pcols)              ! wet deposition of hydrophilic organic carbon
    real(r8) :: ocphidry(pcols)              ! dry deposition of hydrophilic organic carbon
    real(r8) :: ocphodry(pcols)              ! dry deposition of hydrophobic organic carbon
    real(r8) :: dstwet1 (pcols)              ! wet deposition of dust (bin1)
    real(r8) :: dstdry1 (pcols)              ! dry deposition of dust (bin1)
    real(r8) :: dstwet2 (pcols)              ! wet deposition of dust (bin2)
    real(r8) :: dstdry2 (pcols)              ! dry deposition of dust (bin2)
    real(r8) :: dstwet3 (pcols)              ! wet deposition of dust (bin3)
    real(r8) :: dstdry3 (pcols)              ! dry deposition of dust (bin3)
    real(r8) :: dstwet4 (pcols)              ! wet deposition of dust (bin4)
    real(r8) :: dstdry4 (pcols)              ! dry deposition of dust (bin4)
    real(r8), pointer :: nhx_nitrogen_flx(:) ! nitrogen deposition fluxes (kgN/m2/s)
    real(r8), pointer :: noy_nitrogen_flx(:) ! nitrogen deposition fluxes (kgN/m2/s)
  end type cam_out_t

  !---------------------------------------------------------------------------
  ! This is the merged state of sea-ice, land and ocean surface parameterizations
  !---------------------------------------------------------------------------

  type cam_in_t
    integer  :: lchnk                             ! chunk index
    integer  :: ncol                              ! number of active columns
    real(r8) :: asdir           (pcols)           ! albedo: shortwave, direct
    real(r8) :: asdif           (pcols)           ! albedo: shortwave, diffuse
    real(r8) :: aldir           (pcols)           ! albedo: longwave, direct
    real(r8) :: aldif           (pcols)           ! albedo: longwave, diffuse
    real(r8) :: lwup            (pcols)           ! longwave up radiative flux
    real(r8) :: lhf             (pcols)           ! latent heat flux
    real(r8) :: shf             (pcols)           ! sensible heat flux
    real(r8) :: wsx             (pcols)           ! surface u-stress (N)
    real(r8) :: wsy             (pcols)           ! surface v-stress (N)
    real(r8) :: tref            (pcols)           ! ref height surface air temp
    real(r8) :: qref            (pcols)           ! ref height specific humidity
    real(r8) :: u10             (pcols)           ! 10m wind speed
    real(r8) :: ugust_out       (pcols)           ! gustiness added
    real(r8) :: u10_with_gust   (pcols)           ! 10m wind speed with gusts added
    real(r8) :: ts              (pcols)           ! merged surface temp
    real(r8) :: sst             (pcols)           ! sea surface temp
    real(r8) :: snowhland       (pcols)           ! snow depth (liquid water equivalent) over land
    real(r8) :: snowhice        (pcols)           ! snow depth over ice
    real(r8) :: fco2_lnd        (pcols)           ! co2 flux from lnd
    real(r8) :: fco2_ocn        (pcols)           ! co2 flux from ocn
    real(r8) :: fdms            (pcols)           ! dms flux
    real(r8) :: landfrac        (pcols)           ! land area fraction
    real(r8) :: icefrac         (pcols)           ! sea-ice areal fraction
    real(r8) :: ocnfrac         (pcols)           ! ocean areal fraction
    real(r8) :: cflx            (pcols,PCNST)     ! constituent flux (emissions)
    real(r8) :: ustar           (pcols)           ! atm/ocn saved version of ustar
    real(r8) :: re              (pcols)           ! atm/ocn saved version of re
    real(r8) :: ssq             (pcols)           ! atm/ocn saved version of ssq
    real(r8), pointer, dimension(:)   :: ram1     ! aerodynamical resistance (s/m) (pcols)
    real(r8), pointer, dimension(:)   :: fv       ! friction velocity (m/s) (pcols)
    real(r8), pointer, dimension(:)   :: soilw    ! volumetric soil water (m3/m3)
    real(r8), pointer, dimension(:,:) :: depvel   ! deposition velocities
    real(r8), pointer, dimension(:,:) :: dstflx   ! dust fluxes
    real(r8), pointer, dimension(:,:) :: meganflx ! MEGAN fluxes
    real(r8), pointer, dimension(:,:) :: fireflx  ! wild fire emissions
    real(r8), pointer, dimension(:)   :: fireztop ! wild fire emissions vert distribution top
  end type cam_in_t

contains

  subroutine hub2atm_alloc( cam_in )

    ! Allocate space for the surface to atmosphere data type. And initialize the values.

    use seq_drydep_mod   , only: lnd_drydep, n_drydep
    use shr_megan_mod    , only: shr_megan_mechcomps_n
    use shr_fire_emis_mod, only: shr_fire_emis_mechcomps_n

    type(cam_in_t), pointer :: cam_in(:) ! Merged surface state

    integer c, ierror
    character(*), parameter :: sub = 'hub2atm_alloc'

    if (.not. phys_grid_initialized()) call endrun(sub // ': phys_grid not called yet')
    allocate(cam_in(begchunk:endchunk), stat=ierror)
    if (ierror /= 0) then
      write(iulog, *) sub // ': Allocation error: ', ierror
      call endrun(sub // ': allocation error')
    end if

    do c = begchunk, endchunk
      nullify(cam_in(c)%ram1)
      nullify(cam_in(c)%fv)
      nullify(cam_in(c)%soilw)
      nullify(cam_in(c)%depvel)
      nullify(cam_in(c)%dstflx)
      nullify(cam_in(c)%meganflx)
      nullify(cam_in(c)%fireflx)
      nullify(cam_in(c)%fireztop)
    end do
    do c = begchunk, endchunk
      if (active_Sl_ram1) then
        allocate(cam_in(c)%ram1(pcols), stat=ierror)
        if (ierror /= 0) call endrun(sub // ': allocation error ram1')
      end if
      if (active_Sl_fv) then
        allocate (cam_in(c)%fv(pcols), stat=ierror)
        if (ierror /= 0) call endrun(sub // ': allocation error fv')
      end if
      if (active_Sl_soilw) then
        allocate(cam_in(c)%soilw(pcols), stat=ierror)
        if (ierror /= 0) call endrun(sub // ': allocation error soilw')
      end if
      if (active_Fall_flxdst1) then
        ! Assume 4 bins from surface model ....
        allocate(cam_in(c)%dstflx(pcols,4), stat=ierror)
        if (ierror /= 0) call endrun(sub // ': allocation error dstflx')
      end if
      if (active_Fall_flxvoc .and. shr_megan_mechcomps_n > 0) then
        allocate(cam_in(c)%meganflx(pcols,shr_megan_mechcomps_n), stat=ierror)
        if (ierror /= 0) call endrun(sub // ': allocation error meganflx')
      end if
    end do

    if (lnd_drydep .and. n_drydep > 0) then
      do c = begchunk, endchunk
        allocate (cam_in(c)%depvel(pcols,n_drydep), stat=ierror)
        if (ierror /= 0) call endrun(sub // ': allocation error depvel')
      end do
    end if

    if (active_Fall_flxfire .and. shr_fire_emis_mechcomps_n > 0) then
      do c = begchunk,endchunk
        allocate(cam_in(c)%fireflx(pcols,shr_fire_emis_mechcomps_n), stat=ierror)
        if (ierror /= 0) call endrun(sub // ': allocation error fireflx')
        allocate(cam_in(c)%fireztop(pcols), stat=ierror)
        if (ierror /= 0) call endrun(sub // ': allocation error fireztop')
      end do
    end if

    do c = begchunk, endchunk
      cam_in(c)%lchnk     = c
      cam_in(c)%ncol      = get_ncols_p(c)
      cam_in(c)%asdir     = 0
      cam_in(c)%asdif     = 0
      cam_in(c)%aldir     = 0
      cam_in(c)%aldif     = 0
      cam_in(c)%lwup      = 0
      cam_in(c)%lhf       = 0
      cam_in(c)%shf       = 0
      cam_in(c)%wsx       = 0
      cam_in(c)%wsy       = 0
      cam_in(c)%tref      = 0
      cam_in(c)%qref      = 0
      cam_in(c)%u10       = 0
      cam_in(c)%ts        = 0
      cam_in(c)%sst       = 0
      cam_in(c)%snowhland = 0
      cam_in(c)%snowhice  = 0
      cam_in(c)%fco2_lnd  = 0
      cam_in(c)%fco2_ocn  = 0
      cam_in(c)%fdms      = 0
      cam_in(c)%landfrac  = posinf
      cam_in(c)%icefrac   = posinf
      cam_in(c)%ocnfrac   = posinf

      if (associated(cam_in(c)%ram1    )) cam_in(c)%ram1     = 0.1_r8
      if (associated(cam_in(c)%fv      )) cam_in(c)%fv       = 0.1_r8
      if (associated(cam_in(c)%soilw   )) cam_in(c)%soilw    = 0
      if (associated(cam_in(c)%dstflx  )) cam_in(c)%dstflx   = 0
      if (associated(cam_in(c)%meganflx)) cam_in(c)%meganflx = 0

      cam_in(c)%cflx  = 0
      cam_in(c)%ustar = 0
      cam_in(c)%re    = 0
      cam_in(c)%ssq   = 0
      if (lnd_drydep .and. n_drydep > 0) then
        cam_in(c)%depvel = 0._r8
      end if
      if (active_Fall_flxfire .and. shr_fire_emis_mechcomps_n > 0) then
        cam_in(c)%fireflx  = 0
        cam_in(c)%fireztop = 0
      end if
    end do

  end subroutine hub2atm_alloc

  subroutine atm2hub_alloc(cam_out)

    ! Allocate space for the atmosphere to surface data type. And initialize the values.

    type(cam_out_t), pointer :: cam_out(:) ! Atmosphere to surface input

    integer c, ierror
    character(*), parameter :: sub = 'atm2hub_alloc'

    if (.not. phys_grid_initialized()) call endrun(sub // ': phys_grid not called yet')
    allocate(cam_out(begchunk:endchunk), stat=ierror)
    if (ierror /= 0) then
      write(iulog, *) sub // ': Allocation error: ', ierror
      call endrun(sub//': allocation error: cam_out')
    end if

    do c = begchunk,endchunk
      cam_out(c)%lchnk    = c
      cam_out(c)%ncol     = get_ncols_p(c)
      cam_out(c)%tbot     = 0
      cam_out(c)%zbot     = 0
      cam_out(c)%topo     = 0
      cam_out(c)%ubot     = 0
      cam_out(c)%vbot     = 0
      cam_out(c)%qbot     = 0
      cam_out(c)%pbot     = 0
      cam_out(c)%rho      = 0
      cam_out(c)%netsw    = 0
      cam_out(c)%flwds    = 0
      cam_out(c)%precsc   = 0
      cam_out(c)%precsl   = 0
      cam_out(c)%precc    = 0
      cam_out(c)%precl    = 0
      cam_out(c)%soll     = 0
      cam_out(c)%sols     = 0
      cam_out(c)%solld    = 0
      cam_out(c)%solsd    = 0
      cam_out(c)%thbot    = 0
      cam_out(c)%co2prog  = 0
      cam_out(c)%co2diag  = 0
      cam_out(c)%psl      = 0
      cam_out(c)%bcphidry = 0
      cam_out(c)%bcphodry = 0
      cam_out(c)%bcphiwet = 0
      cam_out(c)%ocphidry = 0
      cam_out(c)%ocphodry = 0
      cam_out(c)%ocphiwet = 0
      cam_out(c)%dstdry1  = 0
      cam_out(c)%dstwet1  = 0
      cam_out(c)%dstdry2  = 0
      cam_out(c)%dstwet2  = 0
      cam_out(c)%dstdry3  = 0
      cam_out(c)%dstwet3  = 0
      cam_out(c)%dstdry4  = 0
      cam_out(c)%dstwet4  = 0

      nullify(cam_out(c)%nhx_nitrogen_flx)
      nullify(cam_out(c)%noy_nitrogen_flx)

      if (active_Faxa_nhx) then
        allocate(cam_out(c)%nhx_nitrogen_flx(pcols), stat=ierror)
        if (ierror /= 0) call endrun(sub // ': allocation error nhx_nitrogen_flx')
        cam_out(c)%nhx_nitrogen_flx(:) = 0
      end if
      if (active_Faxa_noy) then
        allocate(cam_out(c)%noy_nitrogen_flx(pcols), stat=ierror)
        if (ierror /= 0) call endrun(sub // ': allocation error noy_nitrogen_flx')
        cam_out(c)%noy_nitrogen_flx(:) = 0
      end if
    end do

  end subroutine atm2hub_alloc

  subroutine atm2hub_deallocate(cam_out)

    type(cam_out_t), pointer :: cam_out(:) ! Atmosphere to surface input

    if (associated(cam_out)) then
      deallocate(cam_out)
    end if
    nullify(cam_out)

  end subroutine atm2hub_deallocate

  subroutine hub2atm_deallocate(cam_in)

    type(cam_in_t), pointer :: cam_in(:) ! Atmosphere to surface input

    integer c

    if (associated(cam_in)) then
      do c = begchunk, endchunk
        if (associated(cam_in(c)%ram1)) then
          deallocate(cam_in(c)%ram1)
          nullify(cam_in(c)%ram1)
        end if
        if (associated(cam_in(c)%fv)) then
          deallocate(cam_in(c)%fv)
          nullify(cam_in(c)%fv)
        end if
        if (associated(cam_in(c)%soilw)) then
          deallocate(cam_in(c)%soilw)
          nullify(cam_in(c)%soilw)
        end if
        if (associated(cam_in(c)%dstflx)) then
          deallocate(cam_in(c)%dstflx)
          nullify(cam_in(c)%dstflx)
        end if
        if (associated(cam_in(c)%meganflx)) then
          deallocate(cam_in(c)%meganflx)
          nullify(cam_in(c)%meganflx)
        end if
        if (associated(cam_in(c)%depvel)) then
          deallocate(cam_in(c)%depvel)
          nullify(cam_in(c)%depvel)
        end if
      end do
      deallocate(cam_in)
    end if
    nullify(cam_in)

  end subroutine hub2atm_deallocate

  subroutine cam_export(state, cam_out, pbuf)

    ! Transfer atmospheric fields into necessary surface data structures

    use physics_types , only: physics_state
    use ppgrid        , only: pver
    use cam_history   , only: outfld
    use chem_surfvals , only: chem_surfvals_get
    use co2_cycle     , only: co2_transport, c_i
    use physconst     , only: rair, mwdry, mwco2, gravit
    use constituents  , only: pcnst
    use physics_buffer, only: pbuf_get_index, pbuf_get_field, physics_buffer_desc

    type(physics_state), intent(in) :: state
    type(cam_out_t), intent(inout) :: cam_out
    type(physics_buffer_desc), pointer :: pbuf(:)

    integer i, m, lchnk, ncol
    integer psl_idx
    integer prec_dp_idx, snow_dp_idx, prec_sh_idx, snow_sh_idx
    integer prec_sed_idx,snow_sed_idx,prec_pcw_idx,snow_pcw_idx

    real(r8), pointer :: psl     (:)
    real(r8), pointer :: prec_dp (:)  ! Total precipitation from ZM convection
    real(r8), pointer :: snow_dp (:)  ! snow from ZM convection
    real(r8), pointer :: prec_sh (:)  ! total precipitation from Hack convection
    real(r8), pointer :: snow_sh (:)  ! snow from Hack convection
    real(r8), pointer :: prec_sed(:)  ! total precipitation from ZM convection
    real(r8), pointer :: snow_sed(:)  ! snow from ZM convection
    real(r8), pointer :: prec_pcw(:)  ! total precipitation from Hack convection
    real(r8), pointer :: snow_pcw(:)  ! snow from Hack convection

    lchnk = state%lchnk
    ncol  = state%ncol

    psl_idx = pbuf_get_index('PSL')
    call pbuf_get_field(pbuf, psl_idx, psl)

    prec_dp_idx  = pbuf_get_index('PREC_DP' , errcode=i)
    snow_dp_idx  = pbuf_get_index('SNOW_DP' , errcode=i)
    prec_sh_idx  = pbuf_get_index('PREC_SH' , errcode=i)
    snow_sh_idx  = pbuf_get_index('SNOW_SH' , errcode=i)
    prec_sed_idx = pbuf_get_index('PREC_SED', errcode=i)
    snow_sed_idx = pbuf_get_index('SNOW_SED', errcode=i)
    prec_pcw_idx = pbuf_get_index('PREC_PCW', errcode=i)
    snow_pcw_idx = pbuf_get_index('SNOW_PCW', errcode=i)

    if (prec_dp_idx  > 0) call pbuf_get_field(pbuf, prec_dp_idx , prec_dp )
    if (snow_dp_idx  > 0) call pbuf_get_field(pbuf, snow_dp_idx , snow_dp )
    if (prec_sh_idx  > 0) call pbuf_get_field(pbuf, prec_sh_idx , prec_sh )
    if (snow_sh_idx  > 0) call pbuf_get_field(pbuf, snow_sh_idx , snow_sh )
    if (prec_sed_idx > 0) call pbuf_get_field(pbuf, prec_sed_idx, prec_sed)
    if (snow_sed_idx > 0) call pbuf_get_field(pbuf, snow_sed_idx, snow_sed)
    if (prec_pcw_idx > 0) call pbuf_get_field(pbuf, prec_pcw_idx, prec_pcw)
    if (snow_pcw_idx > 0) call pbuf_get_field(pbuf, snow_pcw_idx, snow_pcw)

    do i = 1, ncol
      cam_out%tbot (i) = state%t   (i,pver)
      cam_out%thbot(i) = state%t   (i,pver) * state%exner(i,pver)
      cam_out%zbot (i) = state%zm  (i,pver)
      cam_out%topo (i) = state%phis(i) / gravit
      cam_out%ubot (i) = state%u   (i,pver)
      cam_out%vbot (i) = state%v   (i,pver)
      cam_out%pbot (i) = state%pmid(i,pver)
      cam_out%psl  (i) = psl(i)
      cam_out%rho  (i) = cam_out%pbot(i) / (rair * cam_out%tbot(i))
    end do
    do m = 1, pcnst
      do i = 1, ncol
        cam_out%qbot(i,m) = state%q(i,pver,m)
      end do
    end do

    cam_out%co2diag(:ncol) = chem_surfvals_get('CO2VMR') * 1.0e+6_r8
    if (co2_transport()) then
      do i = 1, ncol
        cam_out%co2prog(i) = state%q(i,pver,c_i(4)) * 1.0e+6_r8 * mwdry / mwco2
      end do
    end if
    !
    ! Precipitation and snow rates from shallow convection, deep convection and stratiform processes.
    ! Compute total convective and stratiform precipitation and snow rates
    !
    do i = 1, ncol
      cam_out%precc (i) = 0
      cam_out%precl (i) = 0
      cam_out%precsc(i) = 0
      cam_out%precsl(i) = 0
      if (prec_dp_idx  > 0) cam_out%precc (i) = cam_out%precc (i) + prec_dp (i)
      if (prec_sh_idx  > 0) cam_out%precc (i) = cam_out%precc (i) + prec_sh (i)
      if (prec_sed_idx > 0) cam_out%precl (i) = cam_out%precl (i) + prec_sed(i)
      if (prec_pcw_idx > 0) cam_out%precl (i) = cam_out%precl (i) + prec_pcw(i)
      if (snow_dp_idx  > 0) cam_out%precsc(i) = cam_out%precsc(i) + snow_dp (i)
      if (snow_sh_idx  > 0) cam_out%precsc(i) = cam_out%precsc(i) + snow_sh (i)
      if (snow_sed_idx > 0) cam_out%precsl(i) = cam_out%precsl(i) + snow_sed(i)
      if (snow_pcw_idx > 0) cam_out%precsl(i) = cam_out%precsl(i) + snow_pcw(i)

      ! NOTE: These checks should not be necessary if they exist in the parameterizations
      if (cam_out%precc (i) < 0) cam_out%precc (i) = 0
      if (cam_out%precl (i) < 0) cam_out%precl (i) = 0
      if (cam_out%precsc(i) < 0) cam_out%precsc(i) = 0
      if (cam_out%precsl(i) < 0) cam_out%precsl(i) = 0
      if (cam_out%precsc(i) > cam_out%precc(i)) cam_out%precsc(i) = cam_out%precc(i)
      if (cam_out%precsl(i) > cam_out%precl(i)) cam_out%precsl(i) = cam_out%precl(i)
    end do

  end subroutine cam_export

end module camsrfexch
